/**
 * @file Yolo.cpp
 * @author RittoQin (1264145445@qq.com)
 * @brief 
 * @version 1.0
 * @date 2022/4/18
 *
 * @copyright Copyright SCUT RobotLab(c) 2022
 *
 */

#include <iostream>
#include <opencv2/imgproc.hpp>
#include "Yolo.h"

using namespace std;
using namespace cv;

/**
 * 默认构造
 */
Yolo::Yolo() {
    __model_path = param.MODEL;
    this->setupEngine();
}

/**
 * 构造
 * @param path 模型路径
 */
Yolo::Yolo(const std::string &path) {
    __model_path = path;
    this->setupEngine();
}


/**
 * 设置引擎
 */
void Yolo::setupEngine() {
    // 网络参数与阈值写入
    __layers.resize(3);
    __layers[0] = YoloLayer{param.LAYER_NAMES[0], param.STRIDES[0], param.ANCHORS_S32};
    __layers[1] = YoloLayer{param.LAYER_NAMES[1], param.STRIDES[1], param.ANCHORS_S16};
    __layers[2] = YoloLayer{param.LAYER_NAMES[2], param.STRIDES[2], param.ANCHORS_S8};
    // MNN推理设置
    MNN::ScheduleConfig config; // 设置对象
    config.numThread = 4;
    config.type = MNN_FORWARD_CPU;  // 推理设备类型
    // 后端推理设置
    MNN::BackendConfig backend_config;
    backend_config.memory = MNN::BackendConfig::Memory_High;
    backend_config.power = MNN::BackendConfig::Power_High;
    backend_config.precision = MNN::BackendConfig::Precision_Low;
    config.backendConfig = &backend_config; // 设置后端推理设置
    // 设置网络模型 和 Session
    __net = MNN::Interpreter::createFromFile(__model_path.c_str());
    if (__net == nullptr) {
        throw ("\033[1;31m net is a nullptr! No such a model file!\033[0m");
    }
    __session = __net->createSession(config);
    // 随机数生成器
    RNG rng(getTickCount());
    for (int i = 0; i < param.NUM_CLASSES; ++i) {
        Scalar color = Scalar(rng.uniform(0, 256), rng.uniform(0, 256), rng.uniform(0, 256));
        __colors.push_back(color);
    }
}

/**
 * 推理并返回映射至源图像上的物体边界框容器 => 预处理、启动推理、后处理
 * @param src 源图像
 * @return 物体边界框容器
 */
vector<BoxInfo> Yolo::inference(Mat &src) {
    this->preProcess(src);  // 预处理

    __net->runSession(__session);   // 启动推理

    return this->postProcess(__layers, src.size(), param.INPUT_SIZE, param.NUM_CLASSES,
                             param.OBJ_THRESH, param.SCORE_THRESH, param.NMS_THRESH);
}

/**
 * 预处理 => Resize、BGR2RGB、归一化
 * @param src 原图像
 */
void Yolo::preProcess(Mat &src) {
    // 图像预处理
    Mat src_clone;
    resize(src, src_clone, Size(param.INPUT_SIZE, param.INPUT_SIZE), INTER_AREA);
    cvtColor(src_clone, src_clone, COLOR_BGR2RGB);
    src_clone.convertTo(src_clone, CV_32FC3);
    src_clone /= 255.f; // Normalization
    // wrapping input __tensor, convert nhwc format to nchw format
    /**
     * 这里设置了中间变量nhwc_tensor(Tensorflow格式:nhwc)
     * 使用memcpy将图像数据复制到nhwc_tensor中
     * 再使用net的输入节点张量指针的copyFromHostTensor方法，会根据网络将nhwc转为nchw格式
     */
    std::vector<int> dims{1, param.INPUT_SIZE, param.INPUT_SIZE, 3};    // 维度
    MNN::Tensor* nhwc_tensor = MNN::Tensor::create<float>(dims, NULL, MNN::Tensor::TENSORFLOW);
    float* nhwc_data = nhwc_tensor->host<float>();
    int nhwc_size = nhwc_tensor->size();
    memcpy(nhwc_data, src_clone.data, nhwc_size);  // 从内存复制图像数据到中间Tensor中
    __net->getSessionInput(__session, nullptr)->copyFromHostTensor(
            nhwc_tensor);    // 中间Tensor复制到Session输入节点上(会自动转换为onnx的nchw格式)
    tensor_lock.lock();
    delete nhwc_tensor;
    nhwc_tensor = nullptr;
    tensor_lock.unlock();
}

/**
 * 后处理 => 遍历筛选输出张量
 * @param layers 输出层格式
 * @param src_size 源图像尺寸
 * @param input_size 输入尺寸
 * @param num_classes 类别墅
 * @param obj_thresh 物体置信度阈值
 * @param score_thresh 物体类别置信度阈值
 * @param nms_thresh 极大值抑制IOU阈值
 * @return 物体边界框信息容器
 */
std::vector<BoxInfo>
Yolo::postProcess(std::vector<YoloLayer> &layers, Size src_size, int input_size, int num_classes, float obj_thresh,
                  float score_thresh,
                  float nms_thresh) {
    vector<BoxInfo> results;    // 结果容器
    // 遍历3个不同grid cell尺寸的3通道输出层，需要设置中间张量Host
    for (const auto &layer: layers) {
        MNN::Tensor *tensor = __net->getSessionOutput(__session, layer.name.c_str());  // 输出张量
        MNN::Tensor tensor_host(tensor, tensor->getDimensionType());
        tensor->copyToHostTensor(&tensor_host);
        // __tensor shape => Batch-Channel-Height-Width
        int batchs = tensor_host.shape()[0];
        int channels = tensor_host.shape()[1];
        int height = tensor_host.shape()[2];
        int width = tensor_host.shape()[3];
        int pred_items = tensor_host.shape()[4];

        float *data_ptr = tensor_host.host<float>();    // 输出张量数据指针(以行形式村粗)
        // 遍历同批次的不同个输出张量, 加上channels * height * width * pred_items以指向下一个图像输出的张量
        for (int bi = 0; bi < batchs; ++bi) {
            float *batch_ptr = data_ptr + bi * (channels * height * width * pred_items);
            // 遍历同一个输出张量的不同通道，加上 height * width * pred_items 指向该输出张量的下一通道
            for (int ci = 0; ci < channels; ++ci) {
                float *channel_ptr = batch_ptr + ci * (height * width * pred_items);
                // 遍历同一通道的不同行，加上 width * pred_items 指向该通道的下一行
                for (int hi = 0; hi < height; ++hi) {
                    float *height_ptr = channel_ptr + hi * (width * pred_items);
                    // 遍历同一行的不同个数据， 加上pred_items 指向下一个物体边界框
                    for (int wi = 0; wi < width; ++wi) {
                        float *width_ptr = height_ptr + wi * pred_items;    // 某一个物体边界框头指针
                        float *cls_ptr = width_ptr + 5; // 某一个物体边界框类别得分头指针
                        float obj_conf = sigmoid(width_ptr[4]);    // 物体置信度
                        if (obj_conf < obj_thresh) continue;    // 物体置信度小于阈值，循环下一个物体框
                        // 遍历同一个物体边界框的不同类，找出得分最大的一类
                        float max_score = 0;
                        int max_score_id;
                        for (int cls_id = 0; cls_id < num_classes; ++cls_id) {
                            float score = sigmoid(cls_ptr[cls_id]) * obj_conf;  // 物体类别得分
                            if (score > max_score) {
                                max_score = score;
                                max_score_id = cls_id;
                            }
                        }
                        // 如果最大得分大于阈值则存入边界框容器中
                        if (max_score > score_thresh) {
                            // grid cell的物体中心偏移与宽高 映射至 resize后的图像
                            float cx = (sigmoid(width_ptr[0]) * 2.f - 0.5f + wi) * (float) layer.stride;
                            float cy = (sigmoid(width_ptr[1]) * 2.f - 0.5f + hi) * (float) layer.stride;
                            float w = pow(sigmoid(width_ptr[2]) * 2.f, 2) * layer.anchors[ci].width;
                            float h = pow(sigmoid(width_ptr[3]) * 2.f, 2) * layer.anchors[ci].height;

                            BoxInfo box;    // 生成一个物体边界框
                            box.x1 = std::max(0, std::min(input_size, int((cx - w / 2.f))));
                            box.y1 = std::max(0, std::min(input_size, int((cy - h / 2.f))));
                            box.x2 = std::max(0, std::min(input_size, int((cx + w / 2.f))));
                            box.y2 = std::max(0, std::min(input_size, int((cy + h / 2.f))));
                            box.score = max_score;
                            box.label = max_score_id;
                            results.emplace_back(box);
                        }
                    }
                }
            }
        }
    }
    // 极大值抑制
    this->nms(results, nms_thresh);
    // 坐标变换
    this->coordsTransform(results, input_size, src_size);

    return results;
}

/**
 * 极大值抑制
 * @param boxes 物体便捷边界框容器
 * @param nms_thresh 极大值抑制IOU阈值
 */
void Yolo::nms(std::vector<BoxInfo> &boxes, float nms_thresh) {
    // 根据得分排序
    sort(boxes.begin(), boxes.end(), [](BoxInfo a, BoxInfo b) { return a.score > b.score; });
    // 生成存储边界框面积的容器
    vector<float> vArea(boxes.size());
    for (int i = 0; i < int(boxes.size()); ++i) {
        vArea[i] = (boxes.at(i).x2 - boxes.at(i).x1 + 1)
                   * (boxes.at(i).y2 - boxes.at(i).y1 + 1);
    }
    // 遍历外围边界框 (得分较高)
    for (int i = 0; i < int(boxes.size()); ++i) {
        // 遍历内围边界框 (得分较低)
        for (int j = i + 1; j < int(boxes.size());) {
            // 交集正矩形坐标与尺寸
            float xx1 = MAX(boxes[i].x1, boxes[j].x1);  // 交集正矩形左上x坐标--两个匹配正矩形的最大左上x坐标
            float yy1 = MAX(boxes[i].y1, boxes[j].y1);  // 交集正矩形左上y坐标--两个匹配正矩形的最大左上y坐标
            float xx2 = MIN(boxes[i].x2, boxes[j].x2);  // 交集正矩形右下x坐标--两个匹配正矩形的右下最小x坐标
            float yy2 = MIN(boxes[i].y2, boxes[j].y2);  // 交集正矩形右下y坐标--两个匹配正矩形的右下最小y坐标
            float w = MAX(float(0), xx2 - xx1 + 1); // 交集正矩形宽度
            float h = MAX(float(0), yy2 - yy1 + 1); // 交集正矩形高度
            float inter = w * h;    // 交集面积
            float ovr = inter / (vArea[i] + vArea[j] - inter);  // 交集面积 / (并集面积 + 交集面积)
            // 若交并比IOU大于阈值，则将其删除，j数值不变，使其指向下一个元素
            if (ovr >= nms_thresh) {
                boxes.erase(boxes.begin() + j);
                vArea.erase(vArea.begin() + j);
            } else {
                j++;
            }
        }
    }
}

/**
 * 坐标变换 -- 输入图像 => 源图像
 * @param boxes 物体边界框容器
 * @param input_size 输入图像尺寸
 * @param src_size 源图像尺寸
 */
void Yolo::coordsTransform(std::vector<BoxInfo> &boxes, int input_size, cv::Size src_size) {
    // 变换的宽高比例
    float w_ratio = float(src_size.width) / float(input_size);
    float h_ratio = float(src_size.height) / float(input_size);

    for (auto &box: boxes) {
        box.x1 *= w_ratio;
        box.x2 *= w_ratio;
        box.y1 *= h_ratio;
        box.y2 *= h_ratio;
    }
}

/**
 * 绘制物体边界框
 * @param src 源图像
 * @param boxes 边界框容器
 */
void Yolo::drawBox(cv::Mat &src, std::vector<BoxInfo> &boxes) {
    for (const auto &box: boxes) {
        Rect rect(box.x1, box.y1, (box.x2 - box.x1), (box.y2 - box.y1));
        rectangle(src, rect, __colors[box.label], 2);
        int rect_width = sizeof(param.CLASSES[box.label]) / sizeof(char) * 3.4;
        Rect rect1(rect.x, rect.y - 20, rect_width, 20);
        rectangle(src, rect1, __colors[box.label], -1);
        char text[32];
        sprintf(text, "%s %.1f%%", param.CLASSES[box.label].c_str(), box.score * 100);
        putText(src, text, Point(rect1.x, rect1.y + rect1.height / 1.5), FONT_HERSHEY_SIMPLEX, 0.5,
                Scalar(255, 255, 255), 2);
    }
}